/**
 * Software License, Version 1.0
 * 
 * Copyright 2003 The Trustees of Indiana University.  All rights reserved.
 * 
 *
 *Redistribution and use in source and binary forms, with or without 
 *modification, are permitted provided that the following conditions are met:
 *
 *1) All redistributions of source code must retain the above copyright notice,
 * the list of authors in the original source code, this list of conditions and
 * the disclaimer listed in this license;
 *2) All redistributions in binary form must reproduce the above copyright 
 * notice, this list of conditions and the disclaimer listed in this license in
 * the documentation and/or other materials provided with the distribution;
 *3) Any documentation included with all redistributions must include the 
 * following acknowledgement:
 *
 *"This product includes software developed by the Community Grids Lab. For 
 * further information contact the Community Grids Lab at 
 * http://communitygrids.iu.edu/."
 *
 * Alternatively, this acknowledgement may appear in the software itself, and 
 * wherever such third-party acknowledgments normally appear.
 * 
 *4) The name Indiana University or Community Grids Lab or NaradaBrokering, 
 * shall not be used to endorse or promote products derived from this software 
 * without prior written permission from Indiana University.  For written 
 * permission, please contact the Advanced Research and Technology Institute 
 * ("ARTI") at 351 West 10th Street, Indianapolis, Indiana 46202.
 *5) Products derived from this software may not be called NaradaBrokering, 
 * nor may Indiana University or Community Grids Lab or NaradaBrokering appear
 * in their name, without prior written permission of ARTI.
 * 
 *
 * Indiana University provides no reassurances that the source code provided 
 * does not infringe the patent or any other intellectual property rights of 
 * any other entity.  Indiana University disclaims any liability to any 
 * recipient for claims brought by any other entity based on infringement of 
 * intellectual property rights or otherwise.  
 *
 *LICENSEE UNDERSTANDS THAT SOFTWARE IS PROVIDED "AS IS" FOR WHICH NO 
 *WARRANTIES AS TO CAPABILITIES OR ACCURACY ARE MADE. INDIANA UNIVERSITY GIVES
 *NO WARRANTIES AND MAKES NO REPRESENTATION THAT SOFTWARE IS FREE OF 
 *INFRINGEMENT OF THIRD PARTY PATENT, COPYRIGHT, OR OTHER PROPRIETARY RIGHTS. 
 *INDIANA UNIVERSITY MAKES NO WARRANTIES THAT SOFTWARE IS FREE FROM "BUGS", 
 *"VIRUSES", "TROJAN HORSES", "TRAP DOORS", "WORMS", OR OTHER HARMFUL CODE.  
 *LICENSEE ASSUMES THE ENTIRE RISK AS TO THE PERFORMANCE OF SOFTWARE AND/OR 
 *ASSOCIATED MATERIALS, AND TO THE PERFORMANCE AND VALIDITY OF INFORMATION 
 *GENERATED USING SOFTWARE.
 */
package cgl.narada.transport.proxy;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UnsupportedEncodingException;
import java.net.InetAddress;
import java.net.Socket;
import java.net.UnknownHostException;
import java.util.Enumeration;
import java.util.Hashtable;
import java.util.Random;
import java.util.StringTokenizer;
import java.util.Vector;

import cgl.narada.transport.AccessDeniedException;
import cgl.narada.util.Base64;
import cgl.narada.util.logging.Logger;
import cgl.narada.util.logging.LoggerFactory;
import cgl.narada.util.ntlm.NTLMAuthen;

import com.twmacinta.util.MD5;

import digestauth.DigestAuthen;

/**
 * This class implements functionalities to tunnel through proxy.
 * If required, it also does authentication according to proxy specified
 * authentication method. Proxy on Windows uses username and password in
 * the system. In Unix/Linux, proxy can have its own username/password
 * registry.
 * Originally written by John Yin; refactoring by Hongbin Liu
 * 
 * @version 1.1
 */

public class ProxyTunnel {

    final static Logger log = 
	LoggerFactory.getLogger(ProxyTunnel.class.getName());

    private static final String NARADA_USER_AGENT = "Narada SSL/HTTP";
    private static Random s_rand = new Random(System.currentTimeMillis());
    private ProxyInfo proxyInfo;
    private String proxyHost;
    private int proxyPort;
    private String user;
    private String password;
    private String localhost;
    private String domain;
    private Socket tunnel;
    private int socketTimeout = 4000;

    public ProxyTunnel(ProxyInfo pi) {
	this.proxyInfo = pi;
	proxyHost = proxyInfo.getProxyHost();
	proxyPort = proxyInfo.getProxyPort();
	user = proxyInfo.getUser();
	password = proxyInfo.getPassword();
	domain = proxyInfo.getDomain();

	try {
	    localhost = InetAddress.getLocalHost().getHostName();
	    tunnel = new Socket(proxyHost, proxyPort);
        tunnel.setSoTimeout(socketTimeout);
	} catch (UnknownHostException uhe) { uhe.printStackTrace();}
	catch (IOException ioe) { ioe.printStackTrace();}
    }


    public Socket doProxyConnect(String host, int port)
	throws AccessDeniedException, IOException {
        return doProxyConnect(host,port,null);
    }

    /**
     * ---------------------------------------------------------------------
     * @param authorization An additional authorization header field in the
     *        connection request.
     */
    public  Socket doProxyConnect(String host, int port, String authorization)
	throws AccessDeniedException, IOException {

        OutputStream out = tunnel.getOutputStream();
        InputStream input = tunnel.getInputStream();
        String replyStr = 
	    sendConnectCommand(authorization,host,port,input,out).toLowerCase();

	if (replyStr.indexOf("200") >= 0 &&
	    replyStr.indexOf("connection") >= 0 &&
	    replyStr.indexOf("established") >= 0) {
	    log.info("Proxy tunnel established");
	    return (tunnel);
	} else {
	    if (replyStr.indexOf("407") >= 0 &&
		replyStr.indexOf("proxy") >= 0 &&
		replyStr.indexOf("authentication") >= 0 &&
		replyStr.indexOf("required") >= 0) {
		if (authorization != null) {
		    throw new AccessDeniedException("Authorization failed: "+replyStr);
		} else {
		    tunnel.close();
		    tunnel = new Socket(proxyHost,proxyPort);
            tunnel.setSoTimeout(socketTimeout);
		    log.debug("Do user authentication");
		    doUserAuthentication(host, port,replyStr);
		    return (tunnel);
	       }
	    } else {
		throw new IOException("Unable to tunnel through .  Proxy returns \"" +
				      replyStr + "\"");
            }
	}
	
    }
    
    /**
     * --------------------------------------------------------------------
     */
    private  void doUserAuthentication(String host, int port,String response)
        throws AccessDeniedException, IOException {

	if (user == null || password == null) {
            throw new AccessDeniedException("Authorization to proxy required. "+ 
					    "No username/password set: "+response);
        }

        Vector challenges = new Vector();
        boolean success = false;
        for(int authTypeIndex = response.indexOf("proxy-authenticate: ");
            !success && authTypeIndex >= 0;
            authTypeIndex = response.indexOf("proxy-authenticate: ",authTypeIndex+1)) {

            int newlineIndex = response.indexOf("\n",authTypeIndex);
            if (newlineIndex < 0) newlineIndex = response.length();

            String authenticateStr = response.substring(authTypeIndex,newlineIndex);
            StringTokenizer tokenizer = new StringTokenizer(authenticateStr);
            tokenizer.nextToken();
            String authType = tokenizer.nextToken().toLowerCase();

            String replyField = "Proxy-Authorization: ";

            try {
		if (authType.equals("basic")) {
                        log.debug("Responding to Basic.");
                        replyField += "Basic ";
                        replyField += Base64.encodeString(user+":"+password);
             		log.debug("doUserAuthen " + user + ":" +password +
					   " replyfield " + replyField);
			doProxyConnect(host,port,replyField);
                        success = true;
                } else if (authType.equals("digest")) {
                        log.debug("Responding to Digest.");
                        replyField += constructDigestResponse(authenticateStr,host,port);
                        doProxyConnect(host,port,replyField);
                        success = true;
                } else if (authType.equals("ntlm")) {
                        log.debug("Responding to NTLM");
                        doNTLMAuthentication(host,port);
                        success = true;
                } else {
                        log.debug("Unsupported authentication type: "+authType);
                }
            } catch (Exception e) {
                log.warn("Invalid authentication: "+e.getMessage());
            }
        }

        if (!success)
        {
            throw new AccessDeniedException("Unable to successfully authenticate");
        }
    }

    /**
     * --------------------------------------------------------------------
     */
    private  Socket doNTLMAuthentication(String dstHost,int dstPort) 
	throws IOException {
	
        try {
            NTLMAuthen authen = new NTLMAuthen(domain,localhost,user,password);

            OutputStream output = tunnel.getOutputStream();
            InputStream  input = tunnel.getInputStream();

            // Send Type 1 Message.
            String reply = sendConnectCommand("Proxy-Authorization: "+authen.generateNTLM1Header(),
                                              dstHost,dstPort,input,output);
            String ntlm2 = extractField(reply,"proxy-authenticate: ");
            if (ntlm2.toLowerCase().indexOf("ntlm ") < 0) {
                throw new IOException("Invalid NTLM type 2 message response: "+ntlm2);
            }
            authen.parseNTLM2Nonce(ntlm2);

            reply = sendConnectCommand("Proxy-Authorization: "+ 
				       authen.generateNTLM3Header(),
				       dstHost,dstPort,input,output);
            if(reply.toLowerCase().indexOf("200 connection established") < 0) {
                throw new IOException(reply);
            }
            return tunnel;
        } catch (Exception e) {
	    e.printStackTrace();
            throw new IOException("NTLM Failed: "+e.getMessage());
        }
    }

    /**
     * --------------------------------------------------------------------
     * @param authenticate  
     *         Authentication challenge Proxy-Authenticate: Digest real="blah", nonce="blah",
     * @return The Authorization request header.  Response to Digest authentication.
     */
    private  String constructDigestResponse(String authenticateStr,String host,int port) 
	throws IOException {
        int fieldIndex = authenticateStr.indexOf("digest ");
        authenticateStr = authenticateStr.substring(fieldIndex+"digest ".length(),
                                authenticateStr.length());
        StringTokenizer tokenizer =
            new StringTokenizer(authenticateStr, ",");

        Vector fields = new Vector();
        while(tokenizer.hasMoreElements()) {
            String token = tokenizer.nextToken();
            fields.addElement(token.trim());
        }

        //Proxy-Authorization: Digest username="johnyin", realm="anabas-firewall.anabas-test.anabas.com",
        //qop="auth", algorithm="MD5", uri="/", nonce="b8616a1e6a56aa51389110200000b1cca6c755232e0c0f275248924cbd41",
        //nc=00000001, cnonce="5c518aaf2b8338148956fc2cb9ec565d", response="39c20121cbec79e8ed4309207d9fe967"

        Hashtable params = new Hashtable();
        Enumeration e = fields.elements();
        while(e.hasMoreElements()) {
            String field = (String)e.nextElement();
            int index = field.indexOf("=");
            if (index < 0) continue;
            String param = field.substring(0,index);
            String val   = field.substring(index+1,field.length());

            if (val.startsWith("\"") && val.endsWith("\"")) {
                val = val.substring(1,val.length()-1);
            }

            params.put(param,val);
        }

        // Get the realm
        String realm = (String)params.get("realm");
        if (realm == null) {
            throw new IOException("Digest realm not found in: "+authenticateStr);
        }

        // Get the nonce.
        String nonce = (String)params.get("nonce");
        if (nonce == null) {
            throw new IOException("Digest nonce not found in: "+authenticateStr);
        }

        String qop    = (String)params.get("qop");

        DigestAuthen digest = new DigestAuthen(user, password,"CONNECT","/");
        digest.setAlgorithm("MD5");
        digest.setNonce(nonce);
        digest.setRealm(realm);

        if (qop != null) {
            digest.setQop(qop);
        }

        if (qop != null) {
            byte[] randBytes = new byte[32];
            s_rand.nextBytes(randBytes);
            String cnonce = MD5.asHex(randBytes);
            digest.setCNonce(cnonce);
        }
        return digest.createHeader();
    }

    /**
     * ----------------------------------------------------------------------
     * @param  headerField extra header field ot send.  null means no extra field.
     * @return HTTP response header.
     */
    private String sendConnectCommand(String headerField,
                                     String host,
                                     int    port,
                                     InputStream input,
                                     OutputStream out) throws IOException {
        String msg = "CONNECT " + host + ":" + port + " HTTP/1.0\r\n"
                  + "User-Agent: "+NARADA_USER_AGENT+"\r\n";
        if (headerField != null) {
            msg += headerField  + "\r\n";
        }
        msg += "\r\n";

        byte b[];
        try {
             /*
             * We really do want ASCII7 -- the http protocol doesn't change
             * with locale.
             */
             b = msg.getBytes("ASCII7");
        } catch (UnsupportedEncodingException ignored) {
           /*
            * If ASCII7 isn't there, something serious is wrong, but
            * Paranoia Is Good (tm)
            */
            b = msg.getBytes();
	    ignored.printStackTrace();
        }
        log.debug("Proxy connection request:\n"+msg);
        out.write(b);
        out.flush();

        String replyStr = readHttpReply(input);
        return replyStr;
    }

    /**
     * ----------------------------------------------------------------------
     */
    private  String readHttpReply(InputStream input) throws IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        InputStream    in = input;
        boolean        error = false;
        int newlinesSeen = 0;

        while (newlinesSeen < 2) {
            int i = in.read();
            if (i < 0) {
               throw new IOException("Unexpected EOF from proxy");
            }
            if (i == '\n') {
               ++newlinesSeen;
            } else if (i != '\r') {
                newlinesSeen = 0;
            }
            baos.write(i);
        }

        byte[] msg = baos.toByteArray();
        String header = new String(msg);

	String field  = extractField(header,"content-length:");
	if (field.length() > 0) {
	    try {
		int len = Integer.parseInt(field);
		byte[] body = new byte[len];
		input.read(body);
	    } catch (Exception e) {
		e.printStackTrace();
	    }
	}
        return header;
    }

    /**
     * ----------------------------------------------------------------------
     */
    private String extractField(String header,String fieldName) {
        int index = header.toLowerCase().indexOf(fieldName.toLowerCase());
        if (index < 0)  return "";
        
        int indexEnd = header.indexOf("\n",index);
        if (indexEnd < 0) indexEnd = header.length();
        String field = header.substring(index+fieldName.length(),indexEnd);
        return field.trim();
    }
}
